# -*- coding: utf-8 -*-
"""
#-------------------------------------------------------------------------------
# Nom:         gr2m
#
# Auteur:      Alain Gauthier
#
# Créé le:     22/03/2013
# Copyright:   (c) eaucéa 2014
# Licence:     GNU Public Licence (GPL)

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/

Retranscription du modèle GR2M développé par IRSTEA:
c'est la fonction de transfert pluie-débit du modèle classique GR2M
http://fresno.cemagref.fr/webgr/Modeles/GR2M/fonctionnement_GR2M.htm

Ref :
Mouelhi, S., Michel, C., Perrin, C., Andréassian, V.
Stepwise development of a two-parameter monthly water balance model
(2006) Journal of Hydrology, 318 (1-4), pp. 200-214.

#-------------------------------------------------------------------------------
"""

# modules externes
# ----------------
import numpy as np
import scipy.optimize as opt
import math

# --------------------------------------------------------------------

def gr2m(Ax,Apluie,Aevapo,AniveauIni,Apas=240):
    """
    Fonction de transfert pluie-débit définie selon le modèle GR2M de IRSTEA
    Paramètres de la fonction
    Apas = pas de temps (heures) : 240 valeur par défaut
    AniveauIni = vecteur contenant les niveaux de remplissage initiaux des réservoirs (en %)
    Apluie = série temporelle de pluie donnée au pas de temps Apas
    Aevapo = série temporelle d'ETP donnée au pas de temps Apas

    Paramètre du modèle GR4J :
    Ax[0] = capacité du réservoir de production (X1 > 1mm)
    Ax[1] = paramètre d'échange  (X2 dans [0..50]), par défaut = 0
    Ax[2] = capacité du réservoir de routage    (X3 > 1mm), par défaut = 60 mm

    Résultat : dictionnaire contenant les éléments suivants
    ['Q']   = série temporelle des débits calculés aux pas de temps Apas (en mm)
    ['S']   = série temporelle du volume de remplissage du réservoir de production (en mm)
    ['R']   = série temporelle du volume de remplissage du réservoir de routage (en mm)
    ['ECH'] = échanges d'eau
    ['PN']  = pluie nette

    Ref:
    Perrin, C., Michel, C., André Assian, V.
    Improvement of a parsimonious model for streamflow simulation
    (2003) Journal of Hydrology, 279 (1-4), pp. 275-289.
    """

    # initialisation du résultat
    resultat = dict()
    resultat['Q'] = []
    resultat['S'] = []
    resultat['R'] = []
    resultat['ECH'] = []
    resultat['PN'] = []

    # --------------------------------------------------------------------
    # fonction principale de transfert

    # initialisations ----------------------------------------------------
    nVal = min(Apluie.size,Aevapo.size)
    Qcalc = -9.999*np.ones(nVal)
    Scalc = np.zeros(nVal)
    Rcalc = np.zeros(nVal)
    ECHcalc = np.zeros(nVal)
    PRcalc = np.zeros(nVal)

    # percolation
    if (Apas>=240): PERCFACTOR=1   # Monthly time steps
    if (Apas>=12): PERCFACTOR=2.25 # Daily time steps
    if (Apas<12) : PERCFACTOR=4    # Hourly time steps

    # gestion des données manquantes
    NLACMAX = 5
    nLAC=0

    # paramètres de GR2M (copie locale) : X1,X2 et X3 sont bornées par des valeurs limites
    Lx = Ax.copy()
    Lx[0] = max(Lx[0],1)
    Lx[1] = min(max(Lx[1],0),50)
    Lx[2] = max(Lx[2],1)

    # remplissage initial des réservoirs
    LniveauIni = AniveauIni.copy()
    if(LniveauIni[0]<=0): LniveauIni[0]=0.001
    if(LniveauIni[0]>1) : LniveauIni[0]=1
    if(LniveauIni[1]<=0): LniveauIni[1]=0.001
    if(LniveauIni[1]>1) : LniveauIni[1]=1

    # niveau de remplissage
    Sini=Lx[0]*LniveauIni[0];
    Rini=Lx[2]*LniveauIni[1];

    # boucle sur le pas de temps
    S = Sini
    R = Rini
    for itps in range (nVal):
        P=Apluie[itps]; E=Aevapo[itps]; Q=0
        # gestion des données manquantes
        if((P<0)&(nLAC<=NLACMAX)):
            P=0; nLAC = nLAC +1
        if((E<0)&(nLAC<=NLACMAX)):
            E=0; nLAC = nLAC +1
        if(nLAC>NLACMAX):
            P=0; E=0; nLAC=0
            S = Sini
            R = Rini

        # fonction de transfert au pas de temps itps

        # ......Production.............
        phi = math.tanh(P/Lx[0])
        S1 = (S+Lx[0]*phi)/(1+phi*S/Lx[0])
        P1 = P+S-S1

        # ......Percolation.............
        psi = math.tanh(E/Lx[0])
        S2=S1*(1-psi)/(1+psi*(1-S1/Lx[0]))
        S = S2/pow(1+(S2/Lx[0])**3,1/3.)
        P2=S2-S

        # ......niveau du réservoir de routage et échange souterrain..
        P3=P1+P2
        R1 = R+P3
        ECH=(Lx[1]-1)*R1
        R2=0
        IGF=-R1
        TPr = R1 + ECH
        # Cas d'un réservoir de capacité trop petite
        if (TPr>0):
            R2=TPr
            IGF=ECH
        Q2 = (R2**2)/(R2+Lx[2])
        R= R2-Q2

        #--- résultats
        Qcalc[itps]=Q2
        Scalc[itps]=S
        Rcalc[itps]=R
        ECHcalc[itps]=ECH
        PRcalc[itps]=P3

        #--- gestion des données manquantes (valeurs initiales S et R)----------
        if(nLAC<=NLACMAX):
            Sini=S
            Rini=R

    # fin de boucle sur les pas de temps
    resultat['Q'] = Qcalc
    resultat['S'] = Scalc
    resultat['R'] = Rcalc
    resultat['ECH'] = ECHcalc
    resultat['PN'] = PRcalc

    return resultat

# --------------------------------------------------------------------
# --------------------------------------------------------------------
